from typing import Iterator
import requests
from flask import current_app
import json
import logging

class ChatService:
    @staticmethod
    def stream_chat(
        query: str,
        context: str,
        system_message: str = None,
        model: str = 'openchat'
    ) -> Iterator[str]:
        """流式聊天"""
        try:
            # 构建提示词，添加 Markdown 格式要求
            default_system_message = "你是一个专业的助手。请用 Markdown 格式回复，确保内容结构清晰。"
            if system_message:
                messages = [
                    {"role": "system", "content": f"{system_message}\n请用 Markdown 格式回复。"},
                    {"role": "user", "content": f"基于以下上下文回答问题:\n\n{context}\n\n问题: {query}"}
                ]
            else:
                messages = [
                    {"role": "system", "content": default_system_message},
                    {"role": "user", "content": f"基于以下上下文回答问题:\n\n{context}\n\n问题: {query}"}
                ]
            
            # 记录请求信息
            current_app.logger.debug("Chat request - Query: %s", query.encode('utf-8').decode('utf-8'))
            current_app.logger.debug("Context length: %d chars", len(context))
            
            # 调用 Ollama API
            response = requests.post(
                f"{current_app.config['OLLAMA_BASE_URL']}/api/chat",
                json={
                    "model": model,
                    "messages": messages,
                    "stream": True
                },
                stream=True
            )
            
            if not response.ok:
                current_app.logger.error("Ollama API error: %d - %s", response.status_code, response.text)
                yield json.dumps({"error": "对话服务暂时不可用，请稍后重试"})
                return
            
            for line in response.iter_lines():
                if line:
                    try:
                        data = line.decode('utf-8')
                        current_app.logger.debug("Raw SSE data: %s", data)
                        
                        if data.startswith('data: '):
                            data = data[6:]  # 移除 'data: ' 前缀
                        
                        response_data = json.loads(data)
                        if response_data.get('done', False):
                            continue
                            
                        if 'message' in response_data and 'content' in response_data['message']:
                            content = response_data['message']['content']
                            # 直接返回内容，不再进行额外的 JSON 包装
                            yield content
                            
                    except json.JSONDecodeError as e:
                        current_app.logger.error("JSON parse error: %s, Raw data: %s", str(e), data)
                        continue
                    except Exception as e:
                        current_app.logger.error("Error processing response: %s", str(e))
                        continue
                        
        except requests.exceptions.RequestException as e:
            current_app.logger.error("Network error calling Ollama API: %s", str(e))
            yield json.dumps({"error": "网络连接出现问题，请稍后重试"})
        except Exception as e:
            current_app.logger.error("Unexpected error in chat service: %s", exc_info=True)
            yield json.dumps({"error": "对话服务出现错误，请稍后重试"}) 